import torch
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams
from PIL import Image
import util

# Pass the default decoding hyperparameters of Qwen2-7B-Instruct
# max_tokens is for the maximum length for generation.
# Input the model name or path. Can be GPTQ or AWQ models.
sampling_params = SamplingParams(temperature=0.2,
                                     max_tokens=64,
                                     stop_token_ids=None)
llm = LLM(model="hf-models/Qwen-VL", trust_remote_code=True)


def recognize(image_url, target_text):
    util.download_and_show_image(image_url)
    image = Image.open(util.image_path).convert('RGB')
    # Prepare your prompts
    prompt = f'''
    角色：
    你是一个高效的智能助手，可以精准、广泛地根据提取目标从图片中识别相应信息。
    任务：
    你现在的任务是从图片中提取相关信息，提取目标包括：{target_text}，各提取目标使用|分隔，请自行拆分，可寻找相似目标，但禁止随意识别不相关信息。
    步骤：
    1. 将提取目标 {target_text} 根据“|”符号进行分隔，并分组管理；
    2. 针对每组的目标，分析相近短语，方便后续识别时更加泛化；
    3. 针对每组的目标，按照后续“要求”进行识别，要考虑“步骤2”中的相近短语；
    4. 按照要求进行结果返回。
    要求：
    1. 返回结果使用json格式，每条数据包含一个key-value对，key值为我指定的关键信息，value值为所抽取的结果。
    2. 仅考虑单条结果情况。
    3. 如果认为OCR识别结果中没有关键信息key，则将value赋值为“未找到相关信息”即可。如果图片模糊，请不要随意猜测内容，将value赋值为“无法识别”即可。 
    4. 请只输出json格式的结果，不要包含其它多余文字。
    Picture 1: <img></img>\n
    '''
    inputs = {
        "prompt": prompt,
        "multi_modal_data": {
            "image": image
        },
    }

    # generate outputs
    outputs = llm.generate(inputs, sampling_params)

    # Print the outputs.
    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text
        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

    return outputs